import pandas as pd
import clean_data
import numpy as np
from sklearn.model_selection import StratifiedKFold  # 分层K折交叉验证类似于网格搜索cv=折数
import joblib
from sklearn.model_selection import GridSearchCV
import xgboost as xgb
from sklearn.preprocessing import LabelEncoder  # 文字分类
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import transform_train


def log_train():
    # 逻辑回归
    log_x, y = transform_train.train_do()

    # 处理连续数值类型数据:
    log_x.drop(columns=['Age', 'MonthlyIncome', 'TotalWorkingYears', 'DistanceFromHome'], inplace=True)
    # 标准化
    src = StandardScaler()
    log_x = src.fit_transform(log_x)
    # 标准化模型保存
    joblib.dump(src, '../model/src.pkl')
    model_lr = LogisticRegression(random_state=42, class_weight='balanced')

    # 交叉验证+网格搜索

    param_grid = {
        'C': [0.1, 1, 5, 15, 13, 20],  # 扩大范围，当前 [11~14] 太窄
        'penalty': ['l1', 'l2'],
        'solver': ['liblinear'],  # 先专注 liblinear（更稳）
        'max_iter': [1000]  # 固定足够大的值，避免干扰
    }

    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=22)
    grid_search = GridSearchCV(estimator=model_lr, param_grid=param_grid, cv=cv, n_jobs=-1)
    grid_search.fit(log_x, y)
    joblib.dump(grid_search.best_estimator_, '../model/lr.pkl')


if __name__ == '__main__':
    log_train()
